import os
import numpy as np
import open3d as o3d
from sklearn.linear_model import RANSACRegressor

##### 参数
input_path = "../maps/outdoor_full.pcd"
output_path = "../maps/outdoor_fix.pcd"
bound = [[15, -20, -4], [25, 2, 2]]

# 读取点云
pcd = o3d.io.read_point_cloud(input_path)

# 截取指定区域
bbox = o3d.geometry.AxisAlignedBoundingBox(min_bound=bound[0], max_bound=bound[1])
pcd_crop = pcd.crop(bbox)
o3d.visualization.draw_geometries([pcd_crop], window_name="Cropped Point Cloud")

# 接收命令行输入
cont = input("Apply filtering and plane fitting? (y/N): ")
if cont.lower() != 'y':
    exit()

# 筛滤点云
pcd_crop = pcd_crop.voxel_down_sample(voxel_size=0.02)

# 面拟合
plane_model, inliers = pcd_crop.segment_plane(distance_threshold=0.01,
                                               ransac_n=3,
                                               num_iterations=1000)
[a, b, c, d] = plane_model
print(f"Plane equation: {a:.4f}x + {b:.4f}y + {c:.4f}z + {d:.4f} = 0")

# 计算法向量
normal = np.array([a, b, c])
normal /= np.linalg.norm(normal)

# 计算旋转矩阵
z_axis = np.array([0, 0, 1])
v = np.cross(normal, z_axis)
c = np.dot(normal, z_axis)
s = np.linalg.norm(v)
if s == 0:
    R = np.eye(3)
else:
    vx = np.array([[0, -v[2], v[1]],
                   [v[2], 0, -v[0]],
                   [-v[1], v[0], 0]])
    R = np.eye(3) + vx + vx @ vx * ((1 - c) / (s ** 2))

# 旋转原始点云
pcd.rotate(R, center=(0, 0, 0))

# 保存处理后的点云
o3d.io.write_point_cloud(output_path, pcd)
print(f"Saved leveled point cloud to {output_path}")
